import numpy
import matplotlib.pyplot as plt
from matplotlib import font_manager


class output_graphics:

    def __init__(self, person_state, isolation_state, Frequency_total, N):
        self.Frequency_total = Frequency_total
        self.person_state = person_state
        self.isolation_state = isolation_state
        self.N = N

    '''
        函数功能：循环每天的个体状态和隔离状态记录某类特定的个体序号
    '''

    def estimate_record_current(self, id_record_list, t_record_list, i):
        if i not in id_record_list:
            id_record_list.append(i)
            t_record_list.append(i)
        else:
            pass

    def estimate_record_change(self, id_record_list, t_record_list, last_state, t, i):
        if (i not in id_record_list) and self.person_state[t - 1][i] == last_state:
            id_record_list.append(i)
            t_record_list.append(i)
        else:
            pass


    '''
        函数功能：绘制饼图，描述某种状态的个体去向（状态之间的关系）
    '''
    def relation(self, element_list, lable_list, img_title=''):
        plt.pie(
            x=element_list,
            labels=lable_list,
            autopct='%1.1f%%'
        )
        plt.title(img_title)
        plt.show()

    # def state_change_proportion(self,date, state_star,state_change):
    #


    '''
        功能描述：绘制某特殊日期的的各类群数量（柱形图）
    '''
    def plt_everyday(self, day, person_state_0, person_state_1, person_state_2, person_state_3, isolation_state_1,
                     isolation_state_2):
        my_font = font_manager.FontProperties(family='SimHei', size=24)
        plt.title("第", day, "天", fontproperties=my_font)
        plt.xlabel("人群类型", fontproperties=my_font)
        plt.ylabel("人数", fontproperties=my_font)
        name_list = ['unisolated_susceptible', 'unisolated_lurk', 'unisolated_infected',
                     'unisolated_asymptomatic infected', 'isolation', 'diagnosed']
        num_list = [len(person_state_0), len(person_state_1), len(person_state_2), len(person_state_3),
                    len(isolation_state_1), len(isolation_state_2)]
        plt.bar(range(len(num_list)), num_list, color='rgb', tick_label=name_list)
        plt.show()

    def plt_sum(self, state_0, state_1, state_2, state_3, recover, isolation, besure, death, Total):
        plt.rcParams['font.sans-serif'] = ['SimHei']
        date = list(range(self.Frequency_total))
        y1 = state_0
        y2 = state_1
        y3 = state_2
        y4 = state_3
        y6 = recover
        y7 = isolation
        y8 = besure
        y9 = death
        y10 = Total
        plt.plot(date, y1, label='Not isolated - Susceptible')
        plt.plot(date, y2, label='Not isolated - The lurker')
        plt.plot(date, y3, label='Not isolated - Asymptomatic')
        plt.plot(date, y4, label='Not isolated - Infection')
        plt.plot(date, y6, label='Recover')
        plt.plot(date, y7, label='Isolated')
        plt.plot(date, y8, label='Confirmed person')
        plt.plot(date, y9, label='Death')
        plt.plot(date, y10, label='Total')
        # 生成图表
        # 设置横坐标为day，纵坐标为The number of people，标题为Changes in the number of people
        plt.xlabel('Frequency')
        plt.ylabel('The number of people')
        plt.title('Changes in the number of people')
        # 设置纵坐标刻度
        ordinate_scale = [0, int(self.N / 4), int(self.N / 2), int(self.N / 4 * 3), self.N]
        plt.yticks(ordinate_scale)
        # 设置填充选项：参数分别对应横坐标，纵坐标，纵坐标填充起始值，填充颜色（可以有更多选项）
        # plt.fill_between(year, people, 10, color='green')
        # 显示图表
        plt.legend()
        plt.show()

    '''
        功能描述：绘制某特殊类个体数量的变化图，
                 （该类人数较少利用刻度较小的坐标分析数据）
    '''
    def plt_1(self, state_3, isolation, besure, death, day):
        plt.rcParams['font.sans-serif'] = ['SimHei']
        date = list(range(day))
        y4 = state_3
        y5 = isolation
        y6 = besure
        y7 = death
        plt.plot(date, y4, label='Not isolated - Asymptomatic')
        plt.plot(date, y5, label='Isolated')
        plt.plot(date, y6, label='Confirmed person')
        plt.plot(date, y7, label='Death')
        # 生成图表
        # 设置横坐标为day，纵坐标为The number of people，标题为Changes in the number of people
        plt.xlabel('day')
        plt.ylabel('The number of people')
        plt.title('Changes in the number of people')
        # 设置纵坐标刻度
        ordinate_scale = list(range(0.2 * self.N))
        plt.yticks(ordinate_scale)
        # 设置填充选项：参数分别对应横坐标，纵坐标，纵坐标填充起始值，填充颜色（可以有更多选项）
        # plt.fill_between(year, people, 10, color='green')
        # 显示图表
        plt.legend()
        plt.show()

    '''
        函数功能：绘制某一群特殊个体的每天状态（个体状态和隔离状态）变化        
    '''
    def special_individual_change(self, special_group):
        for special_individual in special_group:
            one_person_state_everyday = []
            one_isolation_state_everyday = []
            for i in range(0, self.Frequency_total):
                one_person_state_everyday.append(self.person_state[i][special_individual])
                one_isolation_state_everyday.append(self.isolation_state[i][special_individual])

            plt.rcParams['font.sans-serif'] = ['SimHei']
            date = list(range(self.Frequency_total))

            y1 = one_person_state_everyday
            y2 = one_isolation_state_everyday

            # 生成图表
            # 设置横坐标为day，纵坐标1为The person state，纵坐标2标题为The person state change everyday

            fig = plt.figure()
            ax1 = fig.add_subplot(111)
            ax1.plot(date, y1, color='r')
            ax1.set_title('The person state change everyday')
            ax1.set_xlabel('Frequency')
            ax1.set_ylabel('The person state', color='r')

            ax2 = ax1.twinx()
            ax2.plot(date, y2, color='b')
            ax2.set_ylabel('The isolation state', color='b')

            # 设置纵坐标刻度
            ordinate_scale1 = [-1, 0, 1, 2, 3, 4, 5, 10]
            ordinate_scale2 = [-1, 0, 1, 2, 3]
            ax1.set_yticks(ordinate_scale1)
            ax2.set_yticks(ordinate_scale2)
            # 设置填充选项：参数分别对应横坐标，纵坐标，纵坐标填充起始值，填充颜色（可以有更多选项）
            # plt.fill_between(year, people, 10, color='green')
            # 显示图表
            plt.show()
